# -*- coding: utf-8 -*-  
'''
训练nlu模型

Created on 2021年10月3日
@author: luoyi
'''
import os
import sys
#    取项目根目录
ROOT_PATH = os.path.abspath(os.path.dirname(__file__)).split('NLU')[0]
ROOT_PATH = ROOT_PATH + "NLU"
sys.path.append(ROOT_PATH)

import utils.conf as conf
from utils.dictionary import Dictionaries
from utils.relationship import Relationship
from data.dataset_baidu_nlu import DatasetQuestions
from models.nlu.nets import NLU


#    初始化字典与意图，实体标注
Dictionaries.instance().load_from_pkl()
Relationship.instance().load_from_file()
print('初始化字典完成。 字典字数：', Dictionaries.instance().size())
print('初始化意图，实体标注 意图数：', Relationship.instance().rel_size(), ' 实体标注数：', Relationship.instance().sot_size())


#    语料库
# 总训练集样本数量： 1452723
# 总验证集样本数量： 181793
#    训练集
batch_size = conf.DATASET_BAIDU.get_batch_size()
steps_per_epoch = 1452723 // batch_size
epochs=conf.DATASET_BAIDU.get_epochs()

dq = DatasetQuestions(dictionary=Dictionaries.instance(), 
                      relationship=Relationship.instance(), 
                      batch_size=batch_size, epochs=epochs, shuffle_buffer_rate=conf.DATASET_BAIDU.get_shuffle_buffer_rate(), 
                      max_sen_len=conf.NLU.get_max_sen_len())
db_train = dq.tensor_db(fpath=conf.DATASET_BAIDU.get_question_train_data_path(), 
                        count=1452723)
db_val = dq.tensor_db(fpath=conf.DATASET_BAIDU.get_question_val_data_path(), 
                      count=128)


#    初始化模型
nlu = NLU(name='nlu',
          #    bert相关配置
          max_sen_len=conf.NLU.get_max_sen_len(), 
          max_sen=conf.BERT.get_max_sen(), 
          vocab_size=Dictionaries.instance().size(), 
          n_block=conf.BERT.get_n_block(), 
          n_head=conf.BERT.get_n_head_attention(), 
          d_model=conf.BERT.get_d_model(), 
          f_model=conf.BERT.get_f_model(), 
          dropout_rate=conf.BERT.get_dropout_rate(), 
          #    意图相关配置
          inform_size=Relationship.instance().rel_size(), 
          #    意图实体标注相关配置
          pos_size=Relationship.instance().sot_size(), 
          #    模型相关配置
          learning_rate=conf.NLU.get_learning_rate(), 
          input_shape=(None, 2, conf.NLU.get_max_sen_len()), 
          loss_lamda_crf=conf.NLU.get_loss_lamda_crf(), 
          loss_lamda_inform=conf.NLU.get_loss_lamda_inform(), 
          auto_assembling=True, is_build=True)
nlu.show_info()

#    喂数据
nlu.train_tensor_db(db_train, db_val, 
                    steps_per_epoch, batch_size, epochs, 
                    auto_save_weights_after_traind=True, auto_save_weights_dir=conf.NLU.get_model_save_weights_path(), 
                    auto_learning_rate_schedule=True, 
                    auto_tensorboard=True, auto_tensorboard_dir=conf.NLU.get_tensorboard_dir_path())

